#include "stdafx.h"
#include "cvxIntrinsicSolver.h"
#include "nvmath.h"
#include "cvxSf.h"
#include "cvxStroke.h"
#include "least_square.h"
#include <conio.h>
#include <time.h>


using cvx::Rgb8UImage;
using cvx::Bw8UImage;
using cvx::Bw32FImage;
using cvx::Bw32SImage;
using cvx::Rgb32FImage;
using nv::vec3f;

//local window size
static const int IS_WSIZE = 3;    //size of the local windows
static const int IS_WSIZE2 = 9;   //WSIZE2 = WSIZE*WSIZE, number of pixels in each window ('n' in the paper)
static const int IS_MSIZE = 12;   //WSIZE2+3, number of rows in the M matrix (see equation 8)
static const int IS_WSIZE_2 = 1;  //WSIZE_2 = WSIZE/2

CvxIntrinsicSolver::CvxIntrinsicSolver()
{

}
CvxIntrinsicSolver::~CvxIntrinsicSolver()
{

}

void CvxIntrinsicSolver::setImage(const Mat &img)
{
	assert(!img.empty());
	assert(img.channels() == 3);
	int w = img.cols;
	int h = img.rows;
	m_orgImg  = img.clone();
	m_edgeMap = cv::Mat::zeros(h, w, CV_8U);
	m_miniShading = cv::Mat::zeros(h, w, CV_32F);
	computeMiniShading();

	m_shading = cv::Mat::zeros(h,w, CV_32F);
	m_albedo  = cv::Mat::zeros(h,w, CV_8UC3);
	m_preShading = cv::Mat::zeros(h,w,CV_32F);
	m_preAlbdeo  = cv::Mat::zeros(h,w,CV_8UC3);
}

void CvxIntrinsicSolver::computeEdgeMap(float threshold)
{
	assert(!m_orgImg.empty());
	assert(!m_edgeMap.empty());

	int w = m_orgImg.cols;
	int h = m_orgImg.rows;

	Rgb8UImage org_img(&m_orgImg);
	Bw8UImage edge_map(&m_edgeMap);
	for (int y = 1; y<h-1; ++y)
	{
		for (int x = 1; x<w-1; ++x)
		{
			int r = org_img[y][x].r;
			int g = org_img[y][x].g;
			int b = org_img[y][x].b;
	//		if(r < 10 || g < 10 || b < 10)
			{
	//			edge_map[y][x] = 0;
	//			continue;
			}

			float v[3];
			float top[3], bot[3], left[3], right[3];
			v[0] = org_img[y][x].r;
			v[1] = org_img[y][x].g;
			v[2] = org_img[y][x].b;

			top[0] = org_img[y-1][x].r;
			top[1] = org_img[y-1][x].g;
			top[2] = org_img[y-1][x].b;

			bot[0] = org_img[y+1][x].r;
			bot[1] = org_img[y+1][x].g;
			bot[2] = org_img[y+1][x].b;

			left[0] = org_img[y][x-1].r;
			left[1] = org_img[y][x-1].g;
			left[2] = org_img[y][x-1].b;

			right[0] = org_img[y][x+1].r;
			right[1] = org_img[y][x+1].g;
			right[2] = org_img[y][x+1].b;

			vec3f vf(v);
			vec3f v1f(top);
			vec3f v2f(bot);
			vec3f v3f(left);
			vec3f v4f(right);

			vf = nv::normalize(vf);
			v1f = nv::normalize(v1f);
			v2f = nv::normalize(v2f);
			v3f = nv::normalize(v3f);
			v4f = nv::normalize(v4f);

			float angle = nv::dot(vf, v1f);
			if (angle < threshold)
			{
				edge_map[y][x] = 255;				
			}
			else 
			{
				angle = nv::dot(vf, v2f);
				if (angle < threshold)
				{
					edge_map[y][x] = 255;
				}
				else
				{
					angle = nv::dot(vf, v3f);
					if (angle < threshold)
					{
						edge_map[y][x] = 255;
					}
					else
					{
						angle = nv::dot(vf, v4f);
						if (angle < threshold)
						{
							edge_map[y][x] = 255;
						}
					}
				}
			}
		}
	}

//		imshow("edge map", m_edgeMap);
//		cvWaitKey(0);
}

void CvxIntrinsicSolver::computeMiniShading()
{
	assert(!m_orgImg.empty());
	assert(m_miniShading.depth() == CV_32F);

	int w = m_orgImg.cols;
	int h = m_orgImg.rows;

	Rgb8UImage org_img(&m_orgImg);
	Bw32FImage min_sha(&m_miniShading);
	for (int y = 0; y<h; ++y)
	{
		for (int x = 0; x<w; ++x)
		{
			int r,g,b;
			r = org_img[y][x].r;
			g = org_img[y][x].g;
			b = org_img[y][x].b;
			r = std::max(r, g);
			r = std::max(r, b);
			min_sha[y][x] = r/255.0;
		}
	}
}
void CvxIntrinsicSolver::undo()
{
	m_preShading.copyTo(m_shading);
	m_preAlbdeo.copyTo(m_albedo);
}

void CvxIntrinsicSolver::savePrevious()
{
	m_shading.copyTo(m_preShading);
	m_albedo.copyTo(m_preAlbdeo);
}
void CvxIntrinsicSolver::saveResult()
{
	Mat saveShading = m_shading.clone();
	saveShading.convertTo(saveShading, CV_8U, 255);

	imwrite("shading.png", saveShading);
	imwrite("albedo.png", m_albedo);
	_cprintf("save file: shading.png, albedo.png\n");
}

void CvxIntrinsicSolver::showResult()
{
	Mat showShading = m_shading.clone();
	showShading.convertTo(showShading, CV_8U, 255.0);

	Mat showAlbedo = m_albedo.clone();
	cv::resize(showShading, showShading, cv::Size(m_orgImg.cols/2, m_orgImg.rows/2));
	cv::resize(showAlbedo, showAlbedo, cv::Size(m_orgImg.cols/2, m_orgImg.rows/2));
	imshow("shading", showShading);
	imshow("albedo", showAlbedo);
	cvWaitKey(5000);	
}

void CvxIntrinsicSolver::saveEdgeMap(const string &fileName)
{
	imwrite(fileName, m_edgeMap);
}


/************************************************************************/
/* 
	CvxLocalSolver
*/
/************************************************************************/

CvxLocalSolver::CvxLocalSolver()
{

}
CvxLocalSolver::~CvxLocalSolver()
{

}

bool CvxLocalSolver::solveRect(const vector<cv::Point> &pathPts, int penWidth, const int deviationThreshold)
{
	//check parameters
	assert(pathPts.size()>2);
	assert(penWidth >= 4);
	
	//adaptive sample
	vector<cv::Rect> rectVec;
	CvxLocalSolver::adaptiveSample(pathPts, penWidth, m_orgImg,  rectVec, deviationThreshold);
	if (rectVec.size() < 2)
	{
		_cprintf("sampled rect area too small\n");
		return false;
	}
	cv::Rect rect = CvxSF::GetBoundingBox(m_orgImg.cols, m_orgImg.rows, penWidth, pathPts, 8);

	//1.index map, unknown map		
	int label;
	CvxStroke aStroke(m_orgImg.cols, m_orgImg.rows);
	aStroke.SetStroke(pathPts, penWidth, label);


	int h = rect.height;
	int w = rect.width;
	//make sure change to local image
	Mat strokeMap = aStroke.m_map(rect).clone();
	Mat localImg = m_orgImg(rect);
	Mat indexMap = cv::Mat::zeros(h, w, CV_32S);
	Mat localShading = cv::Mat::zeros(h, w, CV_32F);
	Mat localEdgeMap = m_edgeMap(rect);
	Mat localAlbedo  = cv::Mat::zeros(h,w, CV_8UC3);
	Mat localMin = m_miniShading(rect);

	for (int i = 0; i<rectVec.size(); ++i)
	{
		rectVec[i].x -= rect.x;
		rectVec[i].y -= rect.y;
	}
	int unKnownNum = 0;
	GetIndexMap(strokeMap, label, localImg, rectVec, indexMap, unKnownNum, localShading);
	_cprintf("unknown num in local solver is %d\n", unKnownNum);

	//2.solve in sampled area
	vector<map<int, double>> leftMatrix;
	vector<double> rightVec;
	GetEquationAtLocalWindow(indexMap, localShading, localEdgeMap, localImg, leftMatrix, rightVec, true);	
	GetEquationAtEdge(indexMap, localShading, localEdgeMap, localImg, leftMatrix, rightVec);
	
	vector<double> result(unKnownNum);
	bool isOk = LeastSquare::solver(leftMatrix, rightVec, true, unKnownNum, &result[0]);

	if (isOk)
	{
		Bw32FImage sha_img(&localShading);
		Bw32SImage idx_map(&indexMap);
		Rgb8UImage alb_img(&localAlbedo);
		Rgb8UImage org_img(&localImg);
		Bw32FImage min_sha(&localMin);	
		for (int y = 0; y<h; ++y)
		{
			for (int x =0; x<w; ++x)
			{
				if (idx_map[y][x] == CIS_SHA_KNOWN)
				{
					if (sha_img[y][x] >= min_sha[y][x])
					{
						alb_img[y][x].r = org_img[y][x].r/sha_img[y][x];
						alb_img[y][x].g = org_img[y][x].g/sha_img[y][x];
						alb_img[y][x].b = org_img[y][x].b/sha_img[y][x];
					}
					else
					{
						sha_img[y][x] = 0.0;
					}

				}
				else 
				{
					if (result[idx_map[y][x]] >= min_sha[y][x])
					{
						sha_img[y][x] = result[idx_map[y][x]];
						alb_img[y][x].r = org_img[y][x].r/sha_img[y][x];
						alb_img[y][x].g = org_img[y][x].g/sha_img[y][x];
						alb_img[y][x].b = org_img[y][x].b/sha_img[y][x];
						idx_map[y][x] = -1;
					}
					//@todo marker error area
					
				}
			}
		}
// 		Mat showShading = localShading.clone();
// 		showShading.convertTo(showShading, CV_8U, 255.0);
// 		namedWindow("shading in local");
// 		namedWindow("albedo in local");
// 		imshow("shading in local", localShading);
// 		imshow("albedo in local", localAlbedo);
// 		cvWaitKey(5000);
// 		destroyWindow("shading in local");
// 		destroyWindow("albedo in local");		
	}
	else
	{
		cout<<"local shading solver error."<<endl;
		return false;
	}

	//update index map
	//5.propagate to nearby area
	propagate();

	//7.update shading and albedo map

	savePrevious();	
	localShading.copyTo(m_shading(rect));
	localAlbedo.copyTo(m_albedo(rect));
	
	//test
// 	localShading.convertTo(localShading, CV_8U, 255.0);
// 	imwrite("local_shading.png", localShading);
// 	imwrite("local_albedo.png", localAlbedo);
	return true;
}

bool CvxLocalSolver::propagate()
{
	return true;
}

void CvxLocalSolver::GetIndexMap(const Mat &strokeMap, int label, const Mat &img, 
								 const vector<cv::Rect> &rectVec, Mat &indexMap, 
								 int &unKnownNum, Mat &localShadingmap)
{
	assert(img.channels() == 3);
	assert(indexMap.depth() == CV_32S);

	int h = strokeMap.rows;
	int w = strokeMap.cols;
	assert(h == img.rows && 
		   h == indexMap.rows &&
		   h == localShadingmap.rows);

	//set sampled area
	Bw32FImage sha_map(&localShadingmap);
	Bw32SImage idx_map(&indexMap);
	Rgb8UImage org_img(&img);
	for (int i = 0; i<rectVec.size(); ++i)
	{
		cv::Rect r = rectVec[i];
		Scalar mean, dev;
		double sha_min = 0;
		Mat m1(img, r);
		cv::meanStdDev(m1, mean, dev);
		sha_min = std::max(mean[0], mean[1]);
		sha_min = std::max(sha_min, mean[2]);
		sha_min /= 255.0;
		sha_min += (1.0 - sha_min)/2;

		Scalar avg_color;
		avg_color[0] = mean[0]/sha_min;
		avg_color[1] = mean[1]/sha_min;
		avg_color[2] = mean[2]/sha_min;

		//set shading and index map to -1 as known 
		for (int y = r.y; y<r.y+r.height; ++y)
		{
			for (int x = r.x; x<r.x+r.width; ++x)
			{
				idx_map[y][x] = CIS_SHA_KNOWN;
				float s_r = org_img[y][x].r/avg_color[2];
				float s_g = org_img[y][x].g/avg_color[1];
				float s_b = org_img[y][x].b/avg_color[0];
				float s = std::max(s_r, s_g);
				s = std::max(s, s_b);
				sha_map[y][x] = s;
			}
		}
	}

	//set unknown index
//	Bw8UImage stroke_map(&strokeMap);
	for (int y = 0; y<h; ++y)
	{
		for (int x = 0;x<w; ++x)
		{
			if (idx_map[y][x] != CIS_SHA_KNOWN)
			{
				idx_map[y][x] = unKnownNum++;
			}
		}
	}

}

void CvxLocalSolver::GetEquationAtLocalWindow(const Mat &idxMap, const Mat &shadingMap, 
											   const Mat &edgeMap, const Mat &orgImage, 
											   vector<map<int, double>> &matrix_vec, 
											   vector<double> &right_vec, bool needLeft)
{
	assert(idxMap.depth() == CV_32S);
	assert(shadingMap.depth() == CV_32F);

	int w = orgImage.cols;
	int h = orgImage.rows;
	double epsilon = 0.01;
	double sqrt_epsilon = sqrt(epsilon);

	/* Set M */
	double M[IS_MSIZE][3]; //M matrix (see equation 8)
	M[IS_MSIZE-3][0] = sqrt_epsilon; M[IS_MSIZE-3][1] = 0; M[IS_MSIZE-3][2] = 0;
	M[IS_MSIZE-2][0] = 0; M[IS_MSIZE-2][1] = sqrt_epsilon; M[IS_MSIZE-2][2] = 0;
	M[IS_MSIZE-1][0] = 0; M[IS_MSIZE-1][1] = 0; M[IS_MSIZE-1][2] = sqrt_epsilon;

	Bw8UImage edge_map(&edgeMap);
	Rgb8UImage org_img(&orgImage);
	Bw32FImage sha_map(&shadingMap);
	Bw32SImage idx_map(&idxMap);
	for (int y = IS_WSIZE_2; y<h-IS_WSIZE_2; ++y)
	{
		for (int x = IS_WSIZE_2; x<w-IS_WSIZE_2; ++x)
		{
			//it is not an edge
			if (edge_map[y][x] == 0)
			{
				//fill the M matrix with the image values in the current window
				//Equ 8
				for (int m=-IS_WSIZE_2; m<=IS_WSIZE_2; m++) 
				{
					int yC = y+m;
					for (int n=-IS_WSIZE_2; n<=IS_WSIZE_2; n++)
					{
						int xC = x+n;
						M[(m+IS_WSIZE_2)*IS_WSIZE + n+IS_WSIZE_2][0] = org_img[yC][xC].r;
						M[(m+IS_WSIZE_2)*IS_WSIZE + n+IS_WSIZE_2][1] = org_img[yC][xC].g;
						M[(m+IS_WSIZE_2)*IS_WSIZE + n+IS_WSIZE_2][2] = org_img[yC][xC].b;
					}
				}
				//get Ni
				Mat Mi  = Mat(IS_MSIZE,3,CV_64F, M);
				Mat Mit = Mi.t();
				Mat MitMi = Mit * Mi;
				Mat MitMi_inv = MitMi.inv();
				Mat Ni = Mi * MitMi_inv * Mit;
				Ni = Mat::eye(Ni.rows, Ni.cols, CV_64F) - Ni;	

				double *pNi = (double*)(Ni.data);

				if (needLeft)
				{
					for (int j = 0; j<Ni.rows; ++j)
					{
						map<int, double> imap;
						double rightVal = 0.0;					
						for (int k = 0; k<Ni.cols-3; ++k)
						{
							int niIdx = j * Ni.cols + k;

							//sIdx position in original map
							int sy = k/IS_WSIZE - IS_WSIZE_2 + y;
							int sx = k%IS_WSIZE - IS_WSIZE_2 + x;
							//shading value already known
							if (idx_map[sy][sx] == CIS_SHA_KNOWN)
							{
								rightVal -= pNi[niIdx] * sha_map[sy][sx];								
							}
							else
							{						
								int sIdx = idx_map[sy][sx];
								imap[sIdx] = pNi[niIdx];	
							}
						}
						if (!imap.empty()) //at least one unknown
						{
							matrix_vec.push_back(imap);
							right_vec.push_back(rightVal);
						}
					}
				}
				else  //don't collect matrix_vec to save time
				{
					for (int j = 0; j<Ni.rows; ++j)
					{
						double rightVal = 0.0;
						bool hasUnknown = false;
						for (int k = 0; k<Ni.cols-3; ++k)
						{
							int niIdx = j * Ni.cols + k;

							//sIdx position in original map
							int sy = k/IS_WSIZE - IS_WSIZE_2 + y;
							int sx = k%IS_WSIZE - IS_WSIZE_2 + x;
							//shading value already known
							if (idx_map[sy][sx] == CIS_SHA_KNOWN)
							{
								rightVal -= pNi[niIdx] * sha_map[sy][sx];								
							}
							else
							{	
								hasUnknown = true;								
							}
						}
						if (hasUnknown) //at least one known
						{							
							right_vec.push_back(rightVal);
						}
					}
				}			
			}
		}
	}		
}

void CvxLocalSolver::GetEquationAtEdge(const Mat &idxMap, const Mat &shadingMap, const Mat &edgeMap, 
										const Mat &orgImage, 
									   vector<map<int, double>> &matrix_vec, 
									   vector<double> &right_vec)
{
	assert(idxMap.depth() == CV_32S);
	assert(shadingMap.depth() == CV_32F);

	int w = orgImage.cols;
	int h = orgImage.rows;	

	Bw8UImage edge_map(&edgeMap);
	Rgb8UImage org_img(&orgImage);
	Bw32FImage sha_map(&shadingMap);
	Bw32SImage idx_map(&idxMap);
	for (int y = IS_WSIZE_2; y<h-IS_WSIZE_2; ++y)
	{
		for (int x = IS_WSIZE_2; x<w-IS_WSIZE_2; ++x)
		{
			//it is an edge
			if (edge_map[y][x] == 255)
			{
				//unknown index numbers
				int idxes[9];
				int colors[9];
				int c_value = 0;
				for (int m=-IS_WSIZE_2; m<=IS_WSIZE_2; m++) 
				{
					int yC = y+m;
					for (int n=-IS_WSIZE_2; n<=IS_WSIZE_2; n++)
					{
						int xC = x+n;
						idxes[(m+1)*3+n+1] = idx_map[yC][xC];
						colors[(m+1)*3+n+1] = 0.3 * org_img[yC][xC].r + 0.59 * org_img[yC][xC].g + 0.11 * org_img[yC][xC].b;
					}
				}

				//3 4 5
				{
					map<int, double> imap;
					double rightVal = 0.0;
					//s3
					if (idxes[3] == -1) rightVal -= sha_map[y][x-1] * (colors[4] - colors[5]);					
					else imap[idxes[3]] = colors[4] - colors[5];
					

					//s4
					if (idxes[4] == -1) rightVal -= sha_map[y][x] * (colors[5] - colors[3]);					
					else imap[idxes[4]] = colors[5] - colors[3];
					

					//s5
					if (idxes[5] == -1) rightVal -= sha_map[y][x+1] * (colors[3] - colors[4]);					
					else imap[idxes[5]] = colors[3] - colors[4];
					
					if (!imap.empty())
					{
						matrix_vec.push_back(imap);
						right_vec.push_back(rightVal);
					}									
				}

				//1 4 7
				{
					map<int, double> imap;
					double rightVal = 0.0;

					//s1
					if (idxes[1] == -1) rightVal -= sha_map[y-1][x] * (colors[4] - colors[7]);					
					else imap[idxes[1]] = colors[4] - colors[7];
					

					//s4
					if (idxes[4] == -1) rightVal -= sha_map[y][x] * (colors[7] - colors[1]);					
					else imap[idxes[4]] = colors[7] - colors[1];
					

					//s7
					if (idxes[7] == -1) rightVal -= sha_map[y+1][x] * (colors[1] - colors[4]);					
					else imap[idxes[7]] = colors[1] - colors[4];
					
					if (!imap.empty())
					{
						matrix_vec.push_back(imap);
						right_vec.push_back(rightVal);
					}									
				}
				//0 4 8
				{
					map<int, double> imap;
					double rightVal = 0.0;

					//s0
					if (idxes[0] == -1) rightVal -= sha_map[y-1][x-1] * (colors[4] - colors[8]);					
					else imap[idxes[0]] = colors[4] - colors[8];
					

					//s4
					if (idxes[4] == -1) rightVal -= sha_map[y][x] * (colors[8] - colors[0]);					
					else imap[idxes[4]] = colors[8] - colors[0];
					

					//s8
					if (idxes[8] == -1) rightVal -= sha_map[y+1][x+1] * (colors[0] - colors[4]);					
					else imap[idxes[8]] = colors[0] - colors[4];
					
					if (!imap.empty())
					{
						matrix_vec.push_back(imap);
						right_vec.push_back(rightVal);
					}
				}

				//2 4 6
				{
					map<int, double> imap;
					double rightVal = 0.0;

					//s2
					if (idxes[2] == -1) rightVal -= sha_map[y-1][x+1] * (colors[4] - colors[6]);					
					else imap[idxes[2]] = colors[4] - colors[6];
					

					//s4
					if (idxes[4] == -1) rightVal -= sha_map[y][x] * (colors[6] - colors[2]);					
					else imap[idxes[4]] = colors[6] - colors[2];
					

					//s6
					if (idxes[6] == -1) rightVal -= sha_map[y+1][x-1] * (colors[2] - colors[4]);					
					else imap[idxes[6]] = colors[2] - colors[4];
					
					if (!imap.empty())
					{
						matrix_vec.push_back(imap);
						right_vec.push_back(rightVal);
					}
				}
			}
		}
	}
}


bool CvxLocalSolver::adaptiveSample(const vector<cv::Point> &pathPts, int penWidth, const Mat &img,
						   vector<cv::Rect> &rectVec, const int deviationThreshold)
{
	assert(!pathPts.empty());
	assert(penWidth >= 4);

	int halfPen = penWidth/2;
	//candidate rect in the right position
	vector<cv::Rect> candidateRect;
	cv::Rect curRect( pathPts[0].x - halfPen, pathPts[0].y - halfPen, penWidth, penWidth);
	for (int i = 1; i<pathPts.size(); ++i)
	{
		int x = pathPts[i].x - halfPen;
		int y = pathPts[i].y - halfPen;
		if (abs(x-curRect.x) > penWidth || abs(y - curRect.y) > penWidth)
		{
			cv::Rect r(x, y, penWidth, penWidth);
			curRect = r;
			candidateRect.push_back(r);
		}	
	}

	//get final result from candidate rect
	
	for (int i = 0; i<candidateRect.size(); ++i)
	{
		Scalar mean, stdDev;
		Mat m1(img, candidateRect[i]);
		cv::meanStdDev(m1, mean, stdDev);
		if (stdDev[0] <= deviationThreshold || stdDev[1] <= deviationThreshold || stdDev[2] <= deviationThreshold)
		{
			rectVec.push_back(candidateRect[i]);
		}
		else if (penWidth >= 20)
		{
			//scale to small size
			cv::Rect sr = candidateRect[i];
			sr.x += halfPen/2;
			sr.y += halfPen/2;
			sr.width = halfPen;
			sr.height = halfPen;
			Mat sm(img, sr);
			cv::meanStdDev(sm, mean, stdDev);

			if (stdDev[0] <= deviationThreshold || stdDev[1] <= deviationThreshold || stdDev[2] <= deviationThreshold)
			{
				rectVec.push_back(sr);
			}
		}
	}
	return true;
}



/************************************************************************/
/* 
	CvxLocalRefineSolver
*/
/************************************************************************/

CvxLocalRefineSolver::CvxLocalRefineSolver()
{

}
CvxLocalRefineSolver::~CvxLocalRefineSolver()
{

}

void CvxLocalRefineSolver::setInitShading(const Mat&shading)
{
	assert(shading.depth() == CV_32F);	
	shading.copyTo(m_shading);	

	Rgb8UImage org_img(&m_orgImg);
	Bw32FImage sha_img(&m_shading);
	Bw32FImage min_sha(&m_miniShading);
	Rgb8UImage alb_img(&m_albedo);
	int h = m_orgImg.rows;
	int w = m_orgImg.cols;
	for (int y = 0; y<h; ++y)
	{
		for (int x =0; x<w; ++x)
		{
			if (sha_img[y][x] >= min_sha[y][x])
			{
				alb_img[y][x].r = org_img[y][x].r/sha_img[y][x];
				alb_img[y][x].g = org_img[y][x].g/sha_img[y][x];
				alb_img[y][x].b = org_img[y][x].b/sha_img[y][x];
			}
		}
	}
}

bool CvxLocalRefineSolver::solveRect(const vector<cv::Point> &pathPts, int penWidth)
{
	if (pathPts.size() <= 2)
	{
		_cprintf("path points less than 3\n");
		return false;
	}

	cv::Rect rect = CvxSF::GetBoundingBox(m_orgImg.cols, m_orgImg.rows, penWidth, pathPts, 8);

	//1.index map, unknown map		
	int label;
	CvxStroke aStroke(m_orgImg.cols, m_orgImg.rows);
	aStroke.SetStroke(pathPts, penWidth, label);

	int h = rect.height;
	int w = rect.width;
	//make sure change to local image
	Mat strokeMap = aStroke.m_map(rect).clone();
	Mat localImg = m_orgImg(rect);
	Mat indexMap = cv::Mat::zeros(h, w, CV_32S);
	Mat localShading = m_shading(rect).clone();
	Mat localEdgeMap = m_edgeMap(rect);
	Mat localAlbedo  = cv::Mat::zeros(h,w, CV_8UC3);
	Mat localMin = m_miniShading(rect);
	
 	int unKnownNum = 0;
 	GetIndexMap(strokeMap, 0, indexMap, unKnownNum);	
 	_cprintf("unknown num in local solver is %d\n", unKnownNum);

	if (unKnownNum == 0)
	{
		return false;
	}
 
// 	//2.solve in sampled area
 	vector<map<int, double>> leftMatrix;
 	vector<double> rightVec;
	CvxLocalSolver::GetEquationAtLocalWindow(indexMap, localShading, localEdgeMap, localImg, leftMatrix, rightVec, true);	
 	CvxLocalSolver::GetEquationAtEdge(indexMap, localShading, localEdgeMap, localImg, leftMatrix, rightVec);
// 
 	vector<double> result(unKnownNum);
 	bool isOk = LeastSquare::solver(leftMatrix, rightVec, true, unKnownNum, &result[0]);
	if (isOk)
	{
		Bw32FImage sha_img(&localShading);
		Bw32SImage idx_map(&indexMap);
		Rgb8UImage alb_img(&localAlbedo);
		Rgb8UImage org_img(&localImg);
		Bw32FImage min_sha(&localMin);	
		for (int y = 0; y<h; ++y)
		{
			for (int x =0; x<w; ++x)
			{
				if (idx_map[y][x] == CIS_SHA_KNOWN)
				{
					if (sha_img[y][x] >= min_sha[y][x])
					{
						alb_img[y][x].r = org_img[y][x].r/sha_img[y][x];
						alb_img[y][x].g = org_img[y][x].g/sha_img[y][x];
						alb_img[y][x].b = org_img[y][x].b/sha_img[y][x];
					}
				}
				else 
				{
					if (result[idx_map[y][x]] >= min_sha[y][x])
					{
						sha_img[y][x] = result[idx_map[y][x]];						
					}
					alb_img[y][x].r = org_img[y][x].r/sha_img[y][x];
					alb_img[y][x].g = org_img[y][x].g/sha_img[y][x];
					alb_img[y][x].b = org_img[y][x].b/sha_img[y][x];
					idx_map[y][x] = -1;
				}				
			}
		}
// 		Mat showShading = localShading.clone();
// 		showShading.convertTo(showShading, CV_8U, 255.0);
// 		imshow("shading in local", localShading);
// 		imshow("albedo in local", localAlbedo);		 		
	}
	else
	{
		cout<<"local shading solver error."<<endl;
		return false;
	}

	//update index map
	//5.propagate to nearby area

	//7.update shading and albedo map
	savePrevious();	
	localShading.copyTo(m_shading(rect));
	localAlbedo.copyTo(m_albedo(rect));

	//test
	//after
// 	localShading.convertTo(localShading, CV_8U, 255.0);
// 	imwrite("local_refine_shading.png", localShading);
// 	imwrite("local_refine_albedo.png", localAlbedo);

	//before
// 	localShading = m_preShading(rect).clone();
// 	localShading.convertTo(localShading, CV_8U, 255.0);
// 	imwrite("before_local_refine_shading.png", localShading);
// 	imwrite("before_local_refine_albedo.png", m_preAlbdeo(rect).clone());

	return true;
}

void CvxLocalRefineSolver::GetIndexMap(const Mat &strokeMap, int label, Mat &indexMap, int &unKnownNum)
{
	assert(indexMap.depth() == CV_32S);

	int h = strokeMap.rows;
	int w = strokeMap.cols;
	assert(h == indexMap.rows);

	//set unknown index
	Bw32SImage idx_map(&indexMap);	
	Bw8UImage stroke_map(&strokeMap);
	for (int y = 0; y<h; ++y)
	{
		for (int x = 0;x<w; ++x)
		{
			if (stroke_map[y][x] == label)
			{
				idx_map[y][x] = CIS_SHA_KNOWN;
			}
			else
			{
				idx_map[y][x] = unKnownNum++;
			}
		}
	}



}







